import gradio as gr
import os
import random
import time
from utils.iat_ws import wav2text
from utils.tts_ws import text2wav


# audio to text
def audio_to_text(audio_path):
    if audio_path is None or "":
        return None
    return wav2text(audio_path)


def chat_completions(messages, gr_states, history):
    """
    chat completion

    Parameters:
        messages:
    Returns:
        response: dict, chat api返回的结果
    """
    if not messages:
        return gr_states, history

    max_retry = 5
    retry = 0
    while retry < max_retry:
        try:
            retry += 1
            # !!!需要调用api聊天
            # content = chatapi(message) ...

            # 虚假的ai回答
            fake_content = [
                '故人西辞黄鹤楼',
                '烟花三月下扬州',
                '孤帆远影碧空尽',
                '唯见长江天际流',
            ]
            content = fake_content[random.randint(0, 3)]

            if content:
                gr_states["history"][-1].append(content)
                history.pop()
                history.append(gr_states["history"][-1])
                break
        except Exception as e:
            print(e)
            pass

    if retry == max_retry:
        gr_states["history"][-1].append("Connection Error: 请求失败，请重试")
        history.pop()
        history.append(gr_states["history"][-1])

    return gr_states, history


def process_tts(text):
    """
    text to speech here

    Parameters:
        text: str, 待转换的文本
    Returns:
        path: str, 保存音频的路径
    """

    time_stamp = time.strftime("%Y%m%d-%H%M%S")
    directory = './audio_cache/'

    path = directory + 'audio_' + time_stamp + '.wav'
    text2wav(text, path)
    return path


def get_audio(gr_states, audio):
    """
    在gradio上渲染audio组件, 更新chatbot组件
    """
    response = gr_states["history"][-1][1]
    if response == "Connection Error: 请求失败，请重试" or response is None:
        gr_states["history"].pop()
        return audio
    else:
        audio = process_tts(response)
    return audio


def format_messages(user_message, gr_states, history):
    """prepare the request data [messages] for the chatbot
    Parameters:
        user_message: str, 用户输入的消息
        gr_states: dict, {"system_prompt": str, "hisotry": List, "user_prompt": str}
        history: list, 聊天记录，一个嵌套列表: [["用户消息", "bot回复"],["用户消息", "bot回复"]]
    """
    messages = [
        {
            "role": "system",
            "content": gr_states["system_prompt"],
        },
    ]

    history.append([user_message, None])
    if len(user_message) > 0:
        gr_states["history"].append([user_message])
        for [usr, bot] in history:
            messages.append({
                "role": "user",
                "content": usr
            })
            if bot:
                messages.append({
                    "role": "assistant",
                    "content": bot
                })
        return messages, gr_states, history
    else:
        return None, gr_states, history


def demo():
    with gr.Blocks(gr.themes.Soft()) as demo:
        demo.title = '模拟面试语音演示'
        gr.Markdown('''<center><font size=6>模拟面试语音演示</font></center>''')
        system_prompt = 'CV 算法岗面试'
        gr_states = gr.State({"system_prompt": system_prompt, "history": []})
        messages = gr.State(None)
        with gr.Tab(label=system_prompt):
            # with gr.Row():
            #     role_info_display = gr.Markdown(f''' # {'面试官'}
            #         {system_prompt}
            #     ''')
            with gr.Row():
                with gr.Column(scale=7):
                    with gr.Row():
                        chatbot = gr.Chatbot(label='对话记录', value=[], render_markdown=False, height=500, visible=True)
                    with gr.Row():
                        with gr.Column(scale=1):
                            input_audio = gr.Audio(label="语音输入", sources='microphone', type="filepath")
                        with gr.Column(scale=1):
                            user_prompt = gr.Textbox(label='语音识别', interactive=False, visible=True)
                            send_button = gr.Button("发送")  # 新增发送按钮
                with gr.Column(scale=3):
                    with gr.Row():
                        audio = gr.Audio(label="AI回答", interactive=False, autoplay=True)

        send_button.click(
            format_messages,
            [user_prompt, gr_states, chatbot],
            [messages, gr_states, chatbot]
        ).then(
            chat_completions,
            [messages, gr_states, chatbot],
            [gr_states, chatbot]
        ).then(
            get_audio,
            [gr_states, audio],
            audio
        ).then(
            lambda: ("", None),  # 清空 user_prompt 和 input_audio
            None,
            [user_prompt, input_audio]
        )

        input_audio.change(audio_to_text, input_audio, user_prompt)
    demo.launch()


if __name__ == '__main__':
    tmp_folder_path = 'audio_cache'
    if not os.path.exists(tmp_folder_path):
        os.makedirs(tmp_folder_path)
    demo()
